#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-10-10


"""
【节点名称】：
    RTDETRDetNode_Cuda
【依赖项安装】：
    pip install spirems
    git clone https://github.com/lyuwenyu/RT-DETR.git
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    spirecv_msgs::2DTargets （检测结果）
    sensor_msgs::CompressedImage （可视化结果，需借助可视化工具）
    std_msgs::Boolean （如果输入节点是数据集，则检测完成发布该话题让输入节点继续工作）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号，如果需要长期后台执行，建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
    realtime_det (bool): 是否是实时检测器，设置为True可以降低延迟，但可能会产生丢帧
【节点参数】：
    confidence (float): 目标得分阈值
    nms_thresh (float): NMS后处理参数
    dataset_name (str): 数据集名称
    pt_model (str): 加载模型名称
【备注】：
    无
"""

import threading
from cProfile import label
from queue import Queue
import numpy as np
import time
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
import sys
sys.path.append("/home/jario/deep/RT-DETR/rtdetrv2_pytorch")
import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image, ImageDraw
from src.core import YAMLConfig
import argparse
import os
import platform


dataset_names = {
    "coco_detection": {
        0: "person",
        1: "bicycle",
        2: "car",
        3: "motorcycle",
        4: "airplane",
        5: "bus",
        6: "train",
        7: "truck",
        8: "boat",
        9: "traffic light",
        10: "fire hydrant",
        11: "stop sign",
        12: "parking meter",
        13: "bench",
        14: "bird",
        15: "cat",
        16: "dog",
        17: "horse",
        18: "sheep",
        19: "cow",
        20: "elephant",
        21: "bear",
        22: "zebra",
        23: "giraffe",
        24: "backpack",
        25: "umbrella",
        26: "handbag",
        27: "tie",
        28: "suitcase",
        29: "frisbee",
        30: "skis",
        31: "snowboard",
        32: "sports ball",
        33: "kite",
        34: "baseball bat",
        35: "baseball glove",
        36: "skateboard",
        37: "surfboard",
        38: "tennis racket",
        39: "bottle",
        40: "wine glass",
        41: "cup",
        42: "fork",
        43: "knife",
        44: "spoon",
        45: "bowl",
        46: "banana",
        47: "apple",
        48: "sandwich",
        49: "orange",
        50: "broccoli",
        51: "carrot",
        52: "hot dog",
        53: "pizza",
        54: "donut",
        55: "cake",
        56: "chair",
        57: "couch",
        58: "potted plant",
        59: "bed",
        60: "dining table",
        61: "toilet",
        62: "tv",
        63: "laptop",
        64: "mouse",
        65: "remote",
        66: "keyboard",
        67: "cell phone",
        68: "microwave",
        69: "oven",
        70: "toaster",
        71: "sink",
        72: "refrigerator",
        73: "book",
        74: "clock",
        75: "vase",
        76: "scissors",
        77: "teddy bear",
        78: "hair drier",
        79: "toothbrush"
    }
}


def trans_det_results(labels, boxes, scores, h, w, dataset, confidence):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = int(h)
    sms_results["width"] = int(w)
    sms_results["targets"] = []
    cls = labels.cpu().detach().numpy().astype(np.float64)
    conf = scores.cpu().detach().numpy().astype(np.float64)
    xywh = boxes.cpu().detach().numpy().astype(np.float64)
    for i in range(cls.shape[1]):
        ann = dict()
        ann["category_name"] = dataset_names[dataset][int(cls[0, i])].strip().replace(' ', '_').lower()
        ann["category_id"] = int(cls[0, i])
        ann["score"] = round(conf[0, i], 3)
        if ann["score"] < confidence:
            continue
        ann["bbox"] = [round(j, 3) for j in xywh[0, i, :].tolist()]
        ann["bbox"][2] = ann["bbox"][2] - ann["bbox"][0]
        ann["bbox"][3] = ann["bbox"][3] - ann["bbox"][1]
        sms_results["targets"].append(ann)

    return sms_results


class RTDETRDetNode_Cuda(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.launch_next_emit = self.get_param("launch_next_emit", True)
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.confidence = self.get_param("confidence", 0.001)
        self.imgsz = self.get_param("imgsz", 640)
        self.dataset_name = self.get_param("dataset_name", "coco_detection")
        self.config = self.get_param("config", "/home/jario/deep/RT-DETR/rtdetrv2_pytorch/configs/rtdetrv2/rtdetrv2_r18vd_120e_coco.yml")
        self.weight = self.get_param("weight", "/home/jario/spirecv2_models/RT-DETR/rtdetrv2_r18vd_120e_coco_rerun_48.1.pth")
        self.device = self.get_param("device", "cuda:0")
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        input_url = '/' + job_name + '/sensor/image_raw' \
            if len(self.specified_input_topic) == 0 else self.specified_input_topic

        output_url = '/' + job_name + '/detector/results' \
            if len(self.specified_output_topic) == 0 else self.specified_output_topic

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._show_writer = Publisher(
            '/' + job_name + '/detector/image_results', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        cfg = YAMLConfig(self.config, resume=self.weight)

        checkpoint = torch.load(self.weight, map_location='cpu')
        if 'ema' in checkpoint:
            state = checkpoint['ema']['module']
        else:
            state = checkpoint['model']

        # NOTE load train mode state -> convert to deploy mode
        cfg.model.load_state_dict(state)

        class Model(nn.Module):
            def __init__(self, ) -> None:
                super().__init__()
                self.model = cfg.model.deploy()
                self.postprocessor = cfg.postprocessor.deploy()

            def forward(self, images, orig_target_sizes):
                outputs = self.model(images)
                outputs = self.postprocessor(outputs, orig_target_sizes)
                return outputs

        self.model = Model().to(self.device)
        self.transforms = T.Compose([
            T.Resize((self.imgsz, self.imgsz)),
            T.ToTensor(),
        ])
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()
        self._next_writer.kill()
        self._image_reader.join()
        self._result_writer.join()
        self._show_writer.join()
        self._next_writer.join()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break
            t1 = time.time()

            msg, img = msg_dict['msg'], msg_dict['img']

            file_name = msg['file_name'] if 'file_name' in msg else ''
            im_pil = Image.fromarray(img)
            w, h = im_pil.size
            orig_size = torch.tensor([w, h])[None].to(self.device)

            # DO Object Detection
            im_data = self.transforms(im_pil)[None].to(self.device)
            output = self.model(im_data, orig_size)
            labels, boxes, scores = output

            res_msg = trans_det_results(labels, boxes, scores, h, w, self.dataset_name, self.confidence)
            res_msg['file_name'] = file_name
            res_msg['dataset'] = self.dataset_name
            if 'client_id' in msg:
                res_msg['client_id'] = msg['client_id']
            if 'file_name' in msg:
                res_msg['file_name'] = msg['file_name']
            if 'img_id' in msg:
                res_msg['img_id'] = msg['img_id']
            if 'img_total' in msg:
                res_msg['img_total'] = msg['img_total']
            res_msg['time_used'] = time.time() - t1
            self._result_writer.publish(res_msg)

            if 'img_total' in msg and self.launch_next_emit:
                next_msg = def_msg('std_msgs::Boolean')
                next_msg['data'] = True
                self._next_writer.publish(next_msg)

            if self.b_use_shm:
                msg = self._show_writer.cvimg2sms_mem(img)
            msg['spirecv_msgs::2DTargets'] = res_msg
            self._show_writer.publish(msg)
            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = RTDETRDetNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
